from structAIDetect.modeltrain.structinfermodel import StructReconstruct
from structAIDetect.modeltrain.disinfermodel import DisplaceInterpolator
import torch
import os
import matplotlib.pyplot as plt
import argparse
import json


INPUT_LEN = 44
HIDDEN_SIZE = 128
NUMLAYERS = 1
TOTAL_MODAL_NUM = 4
MAX_FREQ = 25


def displace_predict(displace_model, modalinfo_edge, pred_len):
    displace_model.eval()
    pred_result = torch.zeros([modalinfo_edge.shape[0], pred_len, modalinfo_edge.shape[2] - TOTAL_MODAL_NUM])
    h_t = torch.zeros(NUMLAYERS, modalinfo_edge.shape[0], HIDDEN_SIZE)
    c_t = torch.zeros(NUMLAYERS, modalinfo_edge.shape[0], HIDDEN_SIZE)
    next_input = modalinfo_edge[:, [0], TOTAL_MODAL_NUM:]
    for pred_idx in range(pred_len):
        next_input = torch.concat((modalinfo_edge[:, [0], :TOTAL_MODAL_NUM], next_input), axis=2)
        outputs, (h_t, c_t) = displace_model(next_input, (h_t, c_t))
        next_input = outputs.unsqueeze(1)
        pred_result[:, [pred_idx], :] = outputs.unsqueeze(1)
    
    return pred_result, h_t


def structure_infer(model, displace_model, modalinfo_edge, pred_len=10):
    model.eval()
    displacement, displace_hidden = displace_predict(displace_model, modalinfo_edge, pred_len)
    cellstate_predict = model(displace_hidden[0, :, :], modalinfo_edge[:, 0, :4])
    cellstate_predict = cellstate_predict[:, :, :41, :]
    cellstate_predict = torch.argmax(cellstate_predict.squeeze(0), dim=0)
    
    return cellstate_predict


def structure_plot(cellstate_predict):
    cell_coor_y = torch.arange(0, cellstate_predict.shape[0])
    cell_coor_x = torch.arange(0, cellstate_predict.shape[1])
    x, y = torch.meshgrid(cell_coor_x, cell_coor_y)
    x = x.reshape(-1, 1).squeeze(-1)
    y = y.reshape(-1, 1).squeeze(-1)
    cellstate_predict = cellstate_predict.T.reshape(-1, 1).squeeze(-1)
    x_uncracked = x[cellstate_predict == 1]
    y_uncracked = y[cellstate_predict == 1]
    x_cracked = x[cellstate_predict == 0]
    y_cracked = y[cellstate_predict == 0]
    plt.figure()
    ax = plt.subplot()
    plt.scatter(x_uncracked, y_uncracked, color='blue')
    plt.scatter(x_cracked, y_cracked, color='red')
    ax.axis('equal')
    plt.show()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--filename', help='待分析文件的文件名', default='measure.json')
    args = parser.parse_args()
    measure_file = os.path.join(os.getcwd(), 'data', 'measure', args.filename)
    with open(measure_file) as f:
        data = json.load(f)

    width = data['width']
    x_coor = torch.linspace(0, width, 10)
    
    measure_point = data['measure_point']
    weight = (x_coor - measure_point[0]) / (measure_point[1] - measure_point[0])
    weight = weight.unsqueeze(0).expand(4, -1)
    modal_freq = torch.tensor(data['modal_freq'])
    modal_freq = modal_freq.unsqueeze(0).unsqueeze(0)
    modal_shape = torch.tensor(data['modal_shape'])
    modal_shape0 = modal_shape[:, 0].unsqueeze(1).expand(-1, weight.shape[1])
    modal_shape1 = modal_shape[:, 1].unsqueeze(1).expand(-1, weight.shape[1])
    modalshape_edge = torch.lerp(modal_shape0, modal_shape1, weight)
    modalshape_edge = modalshape_edge.T.reshape(1, -1).unsqueeze(0)
    modalinfo_edge = torch.concat((modal_freq / MAX_FREQ, modalshape_edge), dim=2)

    displacement_model = DisplaceInterpolator(input_size=INPUT_LEN, hidden_size=HIDDEN_SIZE, num_layers=NUMLAYERS)
    displacement_model.load_state_dict(torch.load(os.path.join('model', 'displace_predict_para.pth')))

    struct_infer_model = StructReconstruct(displace_hidden_size=HIDDEN_SIZE, freqs_size=4)
    struct_infer_model.load_state_dict(torch.load(os.path.join('model', 'struct_reconstruct_para.pth')))

    cellstate_predict = structure_infer(struct_infer_model, displacement_model, modalinfo_edge)
    structure_plot(cellstate_predict)


if __name__ == "__main__":
    main()